/*
 * Copyright (c) 2021 ExtremeVision Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * 本demo采用YOLO5目标检测器
 *
 */

#ifndef COMMON_DET_INFER_H
#define COMMON_DET_INFER_H
#include <memory>
#include <map>
#include "opencv2/core.hpp"
#include "bmnn_utils.h"
#include "bm_wrapper.hpp"

typedef struct BoxInfo
{
    float x1;
    float y1;
    float x2;
    float y2;
    float score;
    int label;
    std::string name;
} BoxInfo;

class SampleDetector
{
public:
    SampleDetector();
    ~SampleDetector();
    bool Init(const std::string &model_name, const std::string& name_files, float thresh);
    bool UnInit();
    bool ProcessImage(const std::vector<bm_image>& images, std::vector<std::vector<BoxInfo>> &detected_objs);
private:
    float get_aspect_scaled_ratio(int src_w, int src_h, int dst_w, int dst_h, bool *pIsAligWidth);
    int Preprocess(const std::vector<bm_image>& images);
    int Postprocess(const std::vector<bm_image>& images, std::vector<std::vector<BoxInfo>>& boxes);
    void NMS(std::vector<BoxInfo> &objects, float thresh);
    int argmax(float* data, int dsize);
    static float sigmoid(float x);

private:
    std::shared_ptr<BMNNContext> m_bmContext;
    std::shared_ptr<BMNNNetwork> m_bmNetwork;
    std::vector<bm_image> m_resized_imgs;
    std::vector<bm_image> m_converto_imgs;

    int m_class_num = 80; // default is coco names
    std::vector<std::string> m_class_names;
    int m_net_h, m_net_w;
    int m_max_batch = 1;
    int output_num;
    int min_dim;
    bmcv_convert_to_attr converto_attr;


private:
    float m_thresh;
    float m_nms_thresh = 0.6;
    int m_dev_id = 0;
};

#endif
